package com.googlecode.classgrep;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.lang.annotation.Annotation;
import java.lang.reflect.Constructor;
import java.lang.reflect.Method;
import java.net.JarURLConnection;
import java.net.URL;
import java.net.URLDecoder;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

import org.apache.commons.lang.ArrayUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.EmptyVisitor;

import com.googlecode.classgrep.classloader.ClassLoaderInterface;
import com.googlecode.classgrep.classloader.ClassLoaderInterfaceDelegate;
import com.googlecode.classgrep.entity.invoker.LocalVariableSignature;
import com.googlecode.classgrep.grep.annotated.AnnotationGrep;
import com.googlecode.classgrep.grep.annotated.ClassGrep;
import com.googlecode.classgrep.grep.annotated.FieldGrep;
import com.googlecode.classgrep.grep.annotated.MethodGrep;
import com.googlecode.classgrep.info.AnnotationInfo;
import com.googlecode.classgrep.info.ClassInfo;
import com.googlecode.classgrep.info.FieldInfo;
import com.googlecode.classgrep.info.HierarchyInfo;
import com.googlecode.classgrep.info.InvokeInfo;
import com.googlecode.classgrep.info.InvokeInfo.FunctionInfo;
import com.googlecode.classgrep.info.MethodInfo;
import com.googlecode.classgrep.utils.CommonUtils;
import com.googlecode.classgrep.utils.URLUtil;
import com.googlecode.classgrep.visitor.annotated.ClassDefVisitor;
import com.googlecode.classgrep.visitor.invoker.AllMethodLocalVarInfoFetchVisitor;
import com.googlecode.classgrep.visitor.invoker.MethodInvokeGrepVisitor;



public class GrepRobot implements GrepRobotInternal{
	
    private static final Log LOG = LogFactory.getLog(GrepRobot.class);

    private ClassLoaderInterface classLoaderInterface;
    
    private boolean excludeParent;
    
    private Test<String> classNameFilter;
	
    private List<String> annotations;
    
	private GrepRobot(){
		this.annotations = new ArrayList<String>();
	}
	
    /**
     * Avoid making getInstance synchronized.
     * see: http://en.wikipedia.org/wiki/Initialization-on-demand_holder_idiom
     */
    private static class LazyHolder {
    	public static final GrepRobot INSTANCE = new GrepRobot();
    }

    /**
     * Public method to return the single instance of this class.
     */
    @SuppressWarnings("unused")
	private static GrepRobotClient getSingleton(){
    	return LazyHolder.INSTANCE;
    }
    
    
    public static GrepRobotClient getInstance() {
    	GrepRobot grep = new GrepRobot();
    	return grep;
    }


	public void setClassLoader(ClassLoader classLoaderInterface) {
		this.classLoaderInterface = new ClassLoaderInterfaceDelegate(classLoaderInterface);
	}
	

	public void setExcludeParent(boolean excludeParent) {
		this.excludeParent = excludeParent;
		
	}

	public boolean isExcludeParent() {
		return excludeParent;
	}


	public void setClassNameFilter(Test<String> classNameFilter) {
		this.classNameFilter = classNameFilter;
	}
    
	
	public void startup() throws IOException {
		Collection<URL> urls = CommonUtils.getUrls(classLoaderInterface, excludeParent);
		Set<String> protocols = new HashSet<String>(){
			
			private static final long serialVersionUID = 1L;
			{
                add("jar");
            }
        };
        
    	invokeInfoMap = new HashMap<FunctionInfo, List<InvokeInfo>>();
    	for(Method calledFun : calledMethod){
    		invokeInfoMap.put(InvokeInfo.buildCommonFunctionInfo(calledFun), new ArrayList<InvokeInfo>());
    	}
    	
    	for(Constructor<?> calledConstr : calledConstructor){
    		invokeInfoMap.put(InvokeInfo.buildConstructorFunctionInfo(calledConstr), new ArrayList<InvokeInfo>());
    	}
    	
    	Map<Class<?>, Map<Class<?> , HierarchyInfo>> allHierarchyInfoMap = new HashMap<Class<?>, Map<Class<?> , HierarchyInfo>>();
    	for(Class<?> p : parents){
    		allHierarchyInfoMap.put(p, new HashMap<Class<?> , HierarchyInfo>());
    	}
		
		List<String> classNames = new ArrayList<String>();
        for (URL location : urls) {
            try {
                if (protocols.contains(location.getProtocol())) {
                    classNames.addAll(jar(location));
                } else if ("file".equals(location.getProtocol())) {
                    try {
                        // See if it's actually a jar
                        URL jarUrl = new URL("jar", "", location.toExternalForm() + "!/");
                        JarURLConnection juc = (JarURLConnection) jarUrl.openConnection();
                        juc.getJarFile();
                        classNames.addAll(jar(jarUrl));
                    } catch (IOException e) {
                        classNames.addAll(file(location));
                    }
                }
            } catch (Exception e) {
                if (LOG.isErrorEnabled())
                    LOG.error("Unable to read URL [" + location.toExternalForm() + "]", e);
            }
        }
        
        StringBuilder unableLoadClass = new StringBuilder("Unable to read class [");
        for (String className : classNames) {
        	if(LOG.isDebugEnabled()){
        		LOG.debug("Prepare to load : " + className);
        	}
            try {
                if (classNameFilter == null || (classNameFilter != null && classNameFilter.test(className))){
                    readClassDef(className);
                    readClassHierarchyInfo(className, allHierarchyInfoMap);
                }
            } catch (Throwable e) {
                if (LOG.isErrorEnabled()){
                	classGrep.getNotLoadedClass().add(className);
                	unableLoadClass.append("\"").append(className).append("\",");
                }
            }
        }
        if(unableLoadClass.length() != "Unable to read class [".length()){
        	unableLoadClass.deleteCharAt(unableLoadClass.length() - 1).append("]");
            LOG.error(unableLoadClass);	
        }
        
        //buildHierarchyInfo
        for(Entry<Class<?>, Map<Class<?> , HierarchyInfo>> entry : allHierarchyInfoMap.entrySet()){
        	Class<?> top = entry.getKey();
        	hierarchyInfoMap.put(top, buildHierarchyInfo(top, entry.getValue()));
        }
        
        //set original method or constructor to MethodInfo
        for(List<MethodInfo> methodInfos : getAnnotatedMethodInfoMap().values()){
        	for(MethodInfo methodInfo : methodInfos){
        		try {
        			Object original;
        			String className = methodInfo.getDeclaringClass().getName();
        		    if(methodInfo.getName().equals("<init>")){
        		    	original = CommonUtils.reflect2Constructor(className, methodInfo.getDesc(), classLoaderInterface.getClassLoader());
        		    }else if(methodInfo.getName().equals("<cinit>")){
        		    	original = CommonUtils.forName(className, true, classLoaderInterface.getClassLoader());
        		    }else{
        		    	original = CommonUtils.reflect2Method(className, methodInfo.getName(), methodInfo.getDesc(), classLoaderInterface.getClassLoader());
        		    }
        		    Method setOriginal = MethodInfo.class.getDeclaredMethod("setOriginal", Object.class);
        		    CommonUtils.forceInvokeMethod(methodInfo, setOriginal, original);
        		} catch (Exception e) {
        			throw new RuntimeException(e.getMessage(), e);
				}
        	}
        }
	}

	/**
     * 
     * @return
     */
    public ClassLoaderInterface getClassLoaderInterface() {
		return classLoaderInterface;
	}
    

	public void addAnnotation(Class<? extends Annotation> annotation) {
		annotations.add(annotation.getName());
	}
	
	public List<String> getAnnotationCriterion() {
		return annotations;
	}
	
    
    //<<<<<<<<<<<<<<<<<<<<<<<<<ClassGrep<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


	private ClassGrep classGrep = new ClassGrep();


	public ClassGrep getClassGrep() {
		return classGrep;
	}
    
	public Map<String, List<ClassInfo>> getAnnotatedClassMap() {
		return CommonUtils.unmodifiableMap(classGrep.getAnnotatedClassMap());
	}


	public List<String> getNotLoadedClass() {
		return CommonUtils.unmodifiableList(classGrep.getNotLoadedClass());
	}


	public List<ClassInfo> findAnnotatedClasses(Class<? extends Annotation> annotation) {
		return CommonUtils.unmodifiableList(classGrep.findAnnotatedClasses(annotation));
	}
    
    //>>>>>>>>>>>>>>>>>>>>>>>>>ClassGrep>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    
    //<<<<<<<<<<<<<<<<<<<<<<<<<AnnotationGrep<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


	private AnnotationGrep annotationGrep = new AnnotationGrep();
	
	public AnnotationGrep getAnnotationGrep() {
		return annotationGrep;
	}
	
	public Map<String, List<AnnotationInfo>> getAnnotatedAnnotationMap() {
		return CommonUtils.unmodifiableMap(annotationGrep.getAnnotatedAnnotationMap());
	}


	public List<AnnotationInfo> findAnnotatedAnnotations(Class<? extends Annotation> annotation) {
		return CommonUtils.unmodifiableList(annotationGrep.findAnnotatedAnnotations(annotation));
	}
	
	
    //>>>>>>>>>>>>>>>>>>>>>>>>>AnnotationGrep>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	
    
    //<<<<<<<<<<<<<<<<<<<<<<<<<FieldGrep<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

	private FieldGrep fieldGrep = new FieldGrep();

	public FieldGrep getFieldGrep() {
		return fieldGrep;
	}


	public Map<String, List<FieldInfo>> getAnnotatedFieldMap() {
		return CommonUtils.unmodifiableMap(fieldGrep.getAnnotatedFieldMap());
	}


	public List<FieldInfo> findAnnotatedFields(Class<? extends Annotation> annotation) {
		return CommonUtils.unmodifiableList(fieldGrep.findAnnotatedFields(annotation));
	}
	
	
    //>>>>>>>>>>>>>>>>>>>>>>>>>FieldGrep>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	
    
    //<<<<<<<<<<<<<<<<<<<<<<<<<FieldGrep<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

	private MethodGrep methodGrep = new MethodGrep();

	public MethodGrep getMethodGrep() {
		return methodGrep;
	}


	public Map<String, List<MethodInfo>> getAnnotatedMethodInfoMap() {
		return CommonUtils.unmodifiableMap(methodGrep.getAnnotatedMethodInfoMap());
	}


	public List<MethodInfo> findAnnotatedMethods(Class<? extends Annotation> annotation) {
		return CommonUtils.unmodifiableList(methodGrep.findAnnotatedMethods(annotation));
	}
	
	
    //>>>>>>>>>>>>>>>>>>>>>>>>>FieldGrep>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	
	
	
	//<<<<<<<<<<<<<<<<<<<<<<<<<Extract Class>>>>>>>>>>>>>>>>>>>>>>>>>

    private List<String> file(URL location) {
        List<String> classNames = new ArrayList<String>();
        File dir = new File(URLDecoder.decode(location.getPath()));
        if ("META-INF".equals(dir.getName())) {
            dir = dir.getParentFile(); // Scrape "META-INF" off
        }
        if (dir.isDirectory()) {
            scanDir(dir, classNames, "");
        }
        return classNames;
    }
    

    private void scanDir(File dir, List<String> classNames, String packageName) {
        File[] files = dir.listFiles();
        for (File file : files) {
            if (file.isDirectory()) {
                scanDir(file, classNames, packageName + file.getName() + ".");
            } else if (file.getName().endsWith(".class")) {
                String name = file.getName();
                name = name.replaceFirst(".class$", "");
                // Classes packaged in an exploded .war (e.g. in a VFS file system) should not
                // have WEB-INF.classes in their package name.
                classNames.add(StringUtils.removeStart(packageName, "WEB-INF.classes.") + name);
            }
        }
    }

	private List<String> jar(URL location) throws IOException {
        URL url = URLUtil.normalizeToFileProtocol(location);
        if (url != null) {
            InputStream in = url.openStream();
            try {
                JarInputStream jarStream = new JarInputStream(in);
                return jar(jarStream);
            } finally {
                in.close();
            }
        } else if (LOG.isDebugEnabled())
            LOG.debug("Unable to read [" + location.toExternalForm() + "]");
        
        return Collections.emptyList();
    }
	
    private List<String> jar(JarInputStream jarStream) throws IOException {
        List<String> classNames = new ArrayList<String>();

        JarEntry entry;
        while ((entry = jarStream.getNextJarEntry()) != null) {
            if (entry.isDirectory() || !entry.getName().endsWith(".class")) {
                continue;
            }
            String className = entry.getName();
            className = className.replaceFirst(".class$", "");

            //war files are treated as .jar files, so takeout WEB-INF/classes
            className = StringUtils.removeStart(className, "WEB-INF/classes/"); 

            className = className.replace('/', '.');
            classNames.add(className);
        }

        return classNames;
    }
        
    public void readClassDef(String className) {
        if (!className.endsWith(".class")) {
            className = className.replace('.', '/') + ".class";
        }
        try {
            URL resource = classLoaderInterface.getResource(className);
            if (resource != null) {
                InputStream in = resource.openStream();
                try {
                    ClassReader classReader = new ClassReader(in);
                    if(!annotations.isEmpty()){
                    	readClassDef4Annotated(classReader);
                    }
                    if(!calledMethod.isEmpty()){
                    	readClassDef4CalledMethod(classReader);
                    }
                } finally {
                    in.close();
                }
            } else {
                throw new RuntimeException("Could not load " + className);
            }
        } catch (IOException e) {
            throw new RuntimeException("Could not load " + className, e);
        }

    }
    
    private void readClassDef4Annotated(ClassReader read){
    	read.accept(new ClassDefVisitor(this), ClassReader.SKIP_DEBUG);
    }
    
    private void readClassDef4CalledMethod(ClassReader read){
    	final Map<String, List<LocalVariableSignature>> allVarSigns = new HashMap<String, List<LocalVariableSignature>>();
    	ClassVisitor localVarVisitor = new AllMethodLocalVarInfoFetchVisitor(allVarSigns);
    	read.accept(localVarVisitor, ClassReader.SKIP_DEBUG);
    			
    	read.accept(new EmptyVisitor(){
        	
        	Type classType;

    		@Override
    		public void visit(int version, int access, String name,
    				String signature, String superName, String[] interfaces) {
    			classType = Type.getObjectType(name);
    		}

    		@Override
    		public MethodVisitor visitMethod(int access, String name, String desc,
    				String signature, String[] exceptions) {
    			//System.out.println("    " + name + " " + desc);
    			if(name.contains("$")){
    				return super.visitMethod(access, name, name, signature, exceptions);
    			}
    			String key = AllMethodLocalVarInfoFetchVisitor.getKey(name, desc);
    			List<LocalVariableSignature> locVarSignList = allVarSigns.get(key);
    			
    			List<FunctionInfo> funInfos = new ArrayList<FunctionInfo>(calledMethod.size() + calledConstructor.size());
    			for(Method method : calledMethod){
    				funInfos.add(InvokeInfo.buildCommonFunctionInfo(method));
    			}
    			
    			for(Constructor<?> constrc : calledConstructor){
    				funInfos.add(InvokeInfo.buildConstructorFunctionInfo(constrc));
    			}
    			
    			return new MethodInvokeGrepVisitor(access, name, desc, classType, locVarSignList,
    					funInfos, invokeInfoMap, classLoaderInterface.getClassLoader());
    		}
        	
        }, ClassReader.SKIP_FRAMES);
    }
    

    private void readClassHierarchyInfo(String className, Map<Class<?>, Map<Class<?> , HierarchyInfo>> childrenMap){
    	try {
			Class<?> clazz = CommonUtils.forName(className, true, classLoaderInterface.getClassLoader());
			Set<Class<?>> parents = childrenMap.keySet();
			for(Class<?> parent : parents){
				if(parent.isAssignableFrom(clazz)){
					childrenMap.get(parent).put(clazz, new HierarchyInfo(clazz));
				}
			}
		} catch (ClassNotFoundException e) {
			LOG.warn("cannot load class : " + className + " cause by : " + e.getMessage(), e);
		}
    }

    private HierarchyInfo buildHierarchyInfo(Class<?> top, Map<Class<?> , HierarchyInfo> childrenMap){
    	
    	for(Class<?> clazz : childrenMap.keySet()){
    		Class<?> superClazz = clazz.getSuperclass();
    		
    		HierarchyInfo currentHier = childrenMap.get(clazz);
    		
    		HierarchyInfo superHier = childrenMap.get(superClazz);
    		if(superHier != null){
    			superHier.getChildren().add(currentHier);
    		}
    		
    		Class<?>[] interfaces = clazz.getInterfaces();
    		if(ArrayUtils.isNotEmpty(interfaces)){
        		for(int i=0; i<interfaces.length; i++){
        			HierarchyInfo interHier = childrenMap.get(interfaces[i]);
        			if(interHier != null){
        				interHier.getChildren().add(currentHier);
            		}
        		}
    		}
    	}
    	
    	return childrenMap.get(top);
    }
    
	//>>>>>>>>>>>>>>>>>>>>>>>>>Extract Class>>>>>>>>>>>>>>>>>>>>>>>>>

    private List<Method> calledMethod = new ArrayList<Method>();
    
    private Map<FunctionInfo, List<InvokeInfo>> invokeInfoMap;

	@Override
	public void addCalledMethod(Method method) {
		calledMethod.add(method);		
	}

	@Override
	public List<Method> getCalledMethods(){
		return calledMethod;
	}

	@Override
	public List<InvokeInfo> getInvokeInfos(Method method) {
		return invokeInfoMap.get(InvokeInfo.buildCommonFunctionInfo(method));
	}

    private List<Constructor<?>> calledConstructor = new ArrayList<Constructor<?>>();

	@Override
	public void addCalledConstructor(Constructor<?> constr) {
		calledConstructor.add(constr);		
	}

	@Override
	public List<Constructor<?>> getCalledConstructor() {
		return calledConstructor;
	}
	
	@Override
	public List<InvokeInfo> getInvokeInfos(Constructor<?> constr) {
		return invokeInfoMap.get(InvokeInfo.buildConstructorFunctionInfo(constr));
	}

	
	private List<Class<?>> parents = new ArrayList<Class<?>>();
	
	private Map<Class<?>, HierarchyInfo> hierarchyInfoMap = new HashMap<Class<?>, HierarchyInfo>();
	
	
	@Override
	public void addParentClass(Class<?> clazz) {
		parents.add(clazz);
	}


	@Override
	public HierarchyInfo getHierarchyChild(Class<?> clazz) {
		return hierarchyInfoMap.get(clazz);
	}


	@Override
	public List<Class<?>> getAllParentClasses() {
		return parents;
	}

	
	
}
